// AES256CBC 流式加解密
// 兼容 OpenSSL 且大大减少内存消耗
// 解密函数无法预测文件大小

package main

import (
	"io"
	"fmt"
)

type aes256CbcStreamEncrypt struct {
	lastIV []byte
	Key    []byte
	buf    []byte
	r      io.Reader
	length int64
}

func (r *aes256CbcStreamEncrypt) Init(key, iv []byte, reader io.Reader, length int64) {
	r.buf = make([]byte, 16)
	r.lastIV = iv
	r.Key = key
	r.r = reader
	r.length = length
}

func (r *aes256CbcStreamEncrypt) Read(p []byte) (n int, err error) {
	n, err = r.r.Read(r.buf)
	if err != nil {
		fmt.Println(err, r.length)
		if r.length > 0 {
			r.length = 0
			n = 0
		} else {
			return
		}
	}

	cipherData, err := AesCbcNoSaltEncrypt(r.buf[:n], r.Key, r.lastIV)
	if err != nil {
		return 0, err
	}
	copy(p, cipherData)
	copy(r.lastIV, cipherData[:16])

	r.length -= 16
	return 16, err
}

type aes256CbcStreamDecrypt struct {
	lastIV []byte
	Key    []byte
	buf    []byte
	r      io.Reader
	length int64
}

func (r *aes256CbcStreamDecrypt) Init(key, iv []byte, reader io.Reader, length int64) {
	r.buf = make([]byte, 16)
	r.lastIV = iv
	r.Key = key
	r.r = reader
	r.length = length
}

func (r *aes256CbcStreamDecrypt) Read(p []byte) (n int, err error) {
	n, err = r.r.Read(r.buf)
	if err != nil || r.length <= 0 {
		return 0, io.EOF
	}
	if r.length == 16 {
		decrypted, err := AesCbcNoSaltDecrypt(r.buf[:16], r.Key, r.lastIV)
		if err != nil {
			return 0, err
		}
		copy(p, decrypted)
		r.length -= 16
		//fmt.Printf("last 16: %X\n", decrypted)
		return len(decrypted), err
	} else {
		decrypted, err := AesCbcNoSaltDecryptWithoutUnpadding(r.buf[:16], r.Key, r.lastIV)
		copy(r.lastIV, r.buf[:16])
		if err != nil {
			return 0, err
		}
		copy(p, decrypted)
		r.length -= 16
		return 16, err
	}
}
